#!/usr/bin/env python
#-*- coding:UTF-8 -*-
import os
import sys
from collections import OrderedDict
import subprocess
import docker
import tempfile
import signal

from knack import CLI
from knack.commands import CLICommandsLoader, CommandGroup
from knack.arguments import ArgumentsContext
from knack.help import CLIHelp
from knack.help_files import helps
from knack.log import CLILogging, get_logger
from knack.parser import CLICommandParser


QINGTIAN_TOOLS_PATH = "/usr/local/share/qingtian/enclave/"
ENCLAVE_LOG = "/tmp/enclave.log"
ENCLAVE_CLI_BIN = "/usr/local/bin/qingtian/qt-enclave"
LINUXKIT_BIN = QINGTIAN_TOOLS_PATH + "linuxkit-linux-amd64"
VMLINUX_FILE = QINGTIAN_TOOLS_PATH + "vmlinux.bin.bz2"
ENCLAVE_INIT = QINGTIAN_TOOLS_PATH + "init"
ENCLAVE_QTSM_KO = QINGTIAN_TOOLS_PATH + "qtsm.ko"
ENCLAVE_DIR = "/var/run/enclave"
DEBUG_MODE = "--debug-mode"

cli_name = os.path.basename(__file__)
logger = get_logger(cli_name)

helps['enclave'] = """
type: group
short-summary: Enclave life-circle management.
"""

helps['enclave make-img'] = """
type: command
short-summary: make an eif image from a docker image
examples:
    - name: Given docker-uri and eif to make an eif image
      text: {cli_name} enclave make-img --docker-uri [DOCKER-URI] --eif [EIF]
    - name: Make an eif image with private-key and signing-certificate
      text: {cli_name} enclave make-img --docker-uri [DOCKER-URI] --eif [EIF] --private-key [PRIVATE-KEY] --signing-certificate [SIGNING-CERTIFICATE]
""".format(cli_name=cli_name)

helps['enclave start'] = """
type: command
short-summary: start an enclave via an eif image
examples:
    - name: Given an eif image, an unused cid, the number of cpus and memory needed
      text:  {cli_name} enclave start  [--cpus CPUS] [--mem MEM] --eif EIF [--cid CID]
    - name: Start an enclave in debug mode
      text:  {cli_name} enclave start  [--cpus CPUS] [--mem MEM] --eif EIF [--cid CID] --debug-mode
""".format(cli_name=cli_name)

helps['enclave stop'] = """
type: command
short-summary: stop an enclave via the enclave-id
examples:
    - name: Given an enclave-id to stop an enclave
      text:  {cli_name} enclave stop --enclave-id [ENCLAVE-ID]
""".format(cli_name=cli_name)

helps['enclave query'] = """
type: command
short-summary: query an enclave via the enclave-id or query all enclaves
examples:
    - name: Given an enclave-id to query an enclave
      text:  {cli_name} enclave query --enclave-id [ENCLAVE-ID]
    - name: Query all enclaves without enclave-id
      text: {cli_name} enclave query
""".format(cli_name=cli_name)

helps['enclave console'] = """
type: command
short-summary: console an enclave via the enclave-id while debugging
examples:
    - name: Given an enclave-id to console an enclave
      text:  {cli_name} enclave console --enclave-id [ENCLAVE-ID]
""".format(cli_name=cli_name)

def execute_cmd(command):
	with subprocess.Popen(command,shell=True) as proc:
		proc.communicate()
		ret = proc.returncode
		return ret

def pull_image(client, image):
    ret = client.api.images(image)
    if len(ret):
        # Use local image
        logger.debug("use the local image")
        return
    else:
        # Pull image from remote uri
        try:
            logger.debug("pull the image from the remote")
            client.api.pull(image)
        except Exception as e:
            logger.error("pull the image fails")
            if isinstance(e, docker.errors.APIError):
                logger.error(e)
            else:
                logger.exception(e)
            sys.exit(1)

def parse_docker(EnclaveCMDFile, EnclaveENVFile, image):
    try:
        client = docker.from_env()
        pull_image(client, image)
        logger.info("pull the image successfully")
        # You have a local docker image now
        # Try to extract the image infomation
        imageInfo = client.api.inspect_image(image)
        cmdContent = imageInfo.get("Config").get("Cmd")
        entryContent = imageInfo.get("Config").get("Entrypoint")
        envContent = imageInfo.get("Config").get("Env")

        # Write the environment variables to EnclaveENVFile
        with open(EnclaveENVFile, 'w') as envfile:
            logger.debug("parse env from the image")
            for envs in envContent:
                envfile.write(envs + "\n")
        # Write the CMD or Entrypoint to
        with open(EnclaveCMDFile, 'w') as cmdfile:
            #CMD or Entrypoint
            if cmdContent is not None:
                # CMD
                logger.debug("parse cmd from the image")
                for parameters in cmdContent:
                    pars = parameters.split(" ")
                    for par in pars:
                        cmdfile.write(par + "\n")
            else:
                # Entrypoint
                for entrypoints in entryContent:
                    logger.debug("parse entrypoints from the image")
                    entryps = entrypoints.split(" ")
                    for entryp in entryps:
                        cmdfile.write(entryp + "\n")
        logger.info("parse docker successfully")
    except Exception as e:
        logger.error("parse the image fails")
        if isinstance(e, docker.errors.APIError):
            logger.error(e)
        else:
            logger.exception(e)
        sys.exit(1)

def generate_yaml(EnclaveCMDFile, EnclaveENVFile, EnclaveYamlFile, uri):
    parse_docker(EnclaveCMDFile, EnclaveENVFile, uri)
    with open(EnclaveYamlFile, 'w') as yaml:
        yaml.write("init:\n")
        yaml.write("  - " + uri + "\n")
        yaml.write("files:\n")
        yaml.write("  - path: init\n    source: " + ENCLAVE_INIT + "\n    mode: \"0755\"\n")
        yaml.write("  - path: qtsm.ko\n    source: " + ENCLAVE_QTSM_KO + "\n    mode: \"0755\"\n")
        yaml.write("  - path: cmd\n    source: " + EnclaveCMDFile + "\n    mode: \"0755\"\n")
        yaml.write("  - path: env\n    source: " + EnclaveENVFile + "\n    mode: \"0755\"\n")
        yaml.write("  - path: /dev\n    directory: true\n    mode: \"0755\"\n")
        yaml.write("  - path: /run\n    directory: true\n    mode: \"0755\"\n")
        yaml.write("  - path: /sys\n    directory: true\n    mode: \"0755\"\n")
        yaml.write("  - path: /var\n    directory: true\n    mode: \"0755\"\n")
        yaml.write("  - path: /proc\n    directory: true\n    mode: \"0755\"\n")
        yaml.write("  - path: /tmp\n    directory: true\n    mode: \"0755\"\n")
        logger.info("generate the yaml file successfully")

def make_img_handler(docker_uri, eif, private_key, signing_certificate):
    logger.debug("start to make an eif image")
    if (private_key and not signing_certificate) or (not private_key and signing_certificate):
        logger.error("please specify private-key and signing-certificate at the same time")
        sys.exit(1)
    EnclaveCMDFile = tempfile.NamedTemporaryFile('w+t')
    EnclaveENVFile = tempfile.NamedTemporaryFile('w+t')
    EnclaveYamlFile = tempfile.NamedTemporaryFile('w+t')
    INITRD_FILE = tempfile.NamedTemporaryFile('w+t')
    try:
        generate_yaml(EnclaveCMDFile.name, EnclaveENVFile.name, EnclaveYamlFile.name, docker_uri)
        command = "%s build -name %s -format kernel+initrd %s > %s 2>&1" % \
                (LINUXKIT_BIN, INITRD_FILE.name, EnclaveYamlFile.name, ENCLAVE_LOG)
        ret = execute_cmd(command)
        if ret != 0:
            logger.error("generating the initrd file fails")
            err_info_cmd = "cat " + ENCLAVE_LOG
            execute_cmd(err_info_cmd)
            sys.exit(1)
        initrd_file = INITRD_FILE.name + "-initrd.img"
        if private_key and signing_certificate:
            command = "%s build_image -k %s -i %s -o %s -p %s -s %s" % \
                    (ENCLAVE_CLI_BIN, VMLINUX_FILE, initrd_file, eif, private_key, signing_certificate)
        else:
            command = "%s build_image -k %s -i %s -o %s" % \
                    (ENCLAVE_CLI_BIN, VMLINUX_FILE, initrd_file, eif)
        ret = execute_cmd(command)
        os.remove(initrd_file)
        if ret != 0:
            sys.exit(1)
    except Exception as e:
        logger.error("make the image fails")
        logger.exception(e)
        sys.exit(1)
    finally:
        EnclaveCMDFile.close()
        EnclaveENVFile.close()
        EnclaveYamlFile.close()
        INITRD_FILE.close()

def start_enclave_handler(cpus, mem, eif, cid, debug_mode):
    if debug_mode:
        logger.debug("start an enclave in debug mode")
        command = "%s run -c %d -m %d -i %s -s %d -d" % \
                (ENCLAVE_CLI_BIN, cpus, mem, eif, cid)
    else:
        logger.debug("start an enclave")
        command = "%s run -c %d -m %d -i %s -s %d" % \
                (ENCLAVE_CLI_BIN, cpus, mem, eif, cid)
    ret = execute_cmd(command)
    if ret != 0:
        sys.exit(1)

def stop_enclave_handler(enclave_id):
    logger.debug("stop an enclave %d", enclave_id)
    command = "%s destroy -i %d" % (ENCLAVE_CLI_BIN, enclave_id)
    ret = execute_cmd(command)
    if ret != 0:
        sys.exit(1)

def query_enclave_handler(enclave_id):
    if enclave_id is not None:
        logger.debug("query an enclave %d", enclave_id)
        command = "%s describe_single_enclave -i%d" % (ENCLAVE_CLI_BIN, enclave_id)
    else:
        logger.debug("query all enclaves")
        command = "%s describe_all_enclaves" % (ENCLAVE_CLI_BIN)
    ret = execute_cmd(command)
    if ret != 0:
        sys.exit(1)

def console_enclave_handler(enclave_id):
    logger.debug("qt console an enclave %d", enclave_id)
    command = "%s console -i %d" % (ENCLAVE_CLI_BIN, enclave_id)
    ret = execute_cmd(command)
    if ret != 0:
        sys.exit(1)

def signal_handler(signum, frame):
    sys.exit(1)

WELCOME_MESSAGE = r"""
  ____  _          _______ _             
 / __ \(_)        |__   __(_)            
| |  | |_ _ __   __ _| |   _  __ _ _ __  
| |  | | | '_ \ / _` | |  | |/ _` | '_ \ 
| |__| | | | | | (_| | |  | | (_| | | | |
 \___\_\_|_| |_|\__, |_|  |_|\__,_|_| |_|
                 __/ |                   
                |___/                    


Welcome to the cool QingTian new CLI!
"""

PRIVACY_STATEMENT = """
Welcome to QingTian CLI!
---------------------
Use `qt -h` to see available commands.
"""

class QtCLIHelp(CLIHelp):
    def __init__(self, cli_ctx=None):
        super(QtCLIHelp, self).__init__(cli_ctx=cli_ctx,
                        privacy_statement=PRIVACY_STATEMENT,
                        welcome_message=WELCOME_MESSAGE)

class QtCommandsLoader(CLICommandsLoader):

    def load_command_table(self, args):
        with CommandGroup(self, 'enclave', '__main__#{}') as g:
            g.command('make-img', 'make_img_handler')
            g.command('start', 'start_enclave_handler')
            g.command('stop', 'stop_enclave_handler')
            g.command('query', 'query_enclave_handler')
            g.command('console', 'console_enclave_handler')
        return super(QtCommandsLoader, self).load_command_table(args)

    def load_arguments(self, command):
        with ArgumentsContext(self, 'enclave start') as ac:
            ac.argument('cpus', type=int, default=2, required=False)
            ac.argument('mem', type=int, default=1024, required=False)
            ac.argument('eif', type=str, required=True)
            ac.argument('cid', type=int, required=True)
            ac.argument('debug_mode', action='store_true', required=False)
        with ArgumentsContext(self, 'enclave make-img') as ac:
            ac.argument('docker_uri', type=str, required=True)
            ac.argument('eif', type=str, required=True)
            ac.argument('private_key', type=str, required=False)
            ac.argument('signing_certificate', type=str, required=False)
        with ArgumentsContext(self, 'enclave stop') as ac:
            ac.argument('enclave_id', type=int, required=True)
        with ArgumentsContext(self, 'enclave query') as ac:
            ac.argument('enclave_id', type=int, default=None, required=False)
        with ArgumentsContext(self, 'enclave console') as ac:
            ac.argument('enclave_id',type=int, required=True)
        super(QtCommandsLoader, self).load_arguments(command)

signal.signal(signal.SIGINT, signal_handler)
os.makedirs(ENCLAVE_DIR, exist_ok=True)
qtcli = CLI(cli_name=cli_name,
            commands_loader_cls=QtCommandsLoader,
            help_cls=QtCLIHelp)
exit_code = qtcli.invoke(sys.argv[1:])
sys.exit(exit_code)
